import pygame
import sys
import random

pygame.init()

# Screen settings
screen_width = 800
screen_height = 600
cell_size = 20
cell_width = screen_width // cell_size
cell_height = screen_height // cell_size
screen = pygame.display.set_mode((screen_width, screen_height))
pygame.display.set_caption('Snake Game')

# Colors
black = (0, 0, 0)
white = (255, 255, 255)
green = (0, 255, 0)
red = (255, 0, 0)

# Snake settings
snake_speed = 15
snake_pos = [100, 50]
snake_body = [[100, 50], [90, 50], [80, 50]]
direction = 'RIGHT'
change_to = direction
score = 0

# Food settings
food_pos = [random.randrange(1, cell_width) * cell_size, random.randrange(1, cell_height) * cell_size]
food_spawn = True

# Game over function
def game_over():
    font = pygame.font.SysFont('Arial', 30)
    text = font.render(f'Your score is: {score}', True, white)
    text_rect = text.get_rect()
    text_rect.midtop = (screen_width // 2, screen_height // 4)
    screen.fill(black)
    screen.blit(text, text_rect)
    pygame.display.flip()
    pygame.time.wait(5000)
    pygame.quit()
    sys.exit()

# Main game loop
while True:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            sys.exit()
        elif event.type == pygame.KEYDOWN:
            if event.key == pygame.K_UP:
                change_to = 'UP'
            if event.key == pygame.K_DOWN:
                change_to = 'DOWN'
            if event.key == pygame.K_LEFT:
                change_to = 'LEFT'
            if event.key == pygame.K_RIGHT:
                change_to = 'RIGHT'

    # Direction validation
    if change_to == 'UP' and direction != 'DOWN':
        direction = 'UP'
    if change_to == 'DOWN' and direction != 'UP':
        direction = 'DOWN'
    if change_to == 'LEFT' and direction != 'RIGHT':
        direction = 'LEFT'
    if change_to == 'RIGHT' and direction != 'LEFT':
        direction = 'RIGHT'

    # Moving the snake
    if direction == 'UP':
        snake_pos[1] -= cell_size
    if direction == 'DOWN':
        snake_pos[1] += cell_size
    if direction == 'LEFT':
        snake_pos[0] -= cell_size
    if direction == 'RIGHT':
        snake_pos[0] += cell_size

    # Snake body mechanism
    snake_body.insert(0, list(snake_pos))
    if snake_pos[0] == food_pos[0] and snake_pos[1] == food_pos[1]:
        score += 1
        food_spawn = False
    else:
        snake_body.pop()

    if not food_spawn:
        food_pos = [random.randrange(1, cell_width) * cell_size, random.randrange(1, cell_height) * cell_size]
        food_spawn = True

    # Background
    screen.fill(black)

    # Draw snake
    for pos in snake_body:
        pygame.draw.rect(screen, green, pygame.Rect(pos[0], pos[1], cell_size, cell_size))

    # Draw food
    pygame.draw.rect(screen, red, pygame.Rect(food_pos[0], food_pos[1], cell_size, cell_size))

    # Game over conditions
    if snake_pos[0] < 0 or snake_pos[0] > screen_width - cell_size:
        game_over()
    if snake_pos[1] < 0 or snake_pos[1] > screen_height - cell_size:
        game_over()
    for block in snake_body[1:]:
        if snake_pos[0] == block[0] and snake_pos[1] == block[1]:
            game_over()

    # Score display
    font = pygame.font.SysFont('Arial', 20)
    score_text = font.render(f'Score: {score}', True, white)
    screen.blit(score_text, (5, 5))

    pygame.display.update()
    pygame.time.Clock().tick(snake_speed)